import matplotlib.pyplot as plt
import multinetx as mx
import numpy as np

g1 = mx.Graph()
g2 = mx.Graph()

g1.add_nodes_from([0,1,2,3,4])
g1.add_edges_from([(0,1), (0,2), (1,2), (1,3), (2,3), (2,4)])

g2.add_nodes_from([0,1,2,3])
g2.add_edges_from([(0,1), (1,2), (1,3)])

adj_block = mx.lil_matrix(np.zeros((9,9)))
adj_block[0:5, 5:9] = np.matrix([[1, 0, 0, 0], 
                                 [0, 1, 0, 0], 
                                 [0, 1, 0, 0], 
                                 [0, 0, 1, 0], 
                                 [0, 0, 0, 1]])
adj_block += adj_block.T
mg = mx.MultilayerGraph(list_of_layers=[g1,g2],
                        inter_adjacency_matrix=adj_block)

# N = 5
# g1 = mx.generators.erdos_renyi_graph(N,0.5,seed=218)
# g2 = mx.generators.erdos_renyi_graph(N,0.6,seed=211)
# g3 = mx.generators.erdos_renyi_graph(N,0.7,seed=208)
# adj_block = mx.lil_matrix(np.zeros((N*3,N*3)))
# adj_block[0:  N,  N:2*N] = np.identity(N)    # L_12
# adj_block[0:  N,2*N:3*N] = np.identity(N)    # L_13
# adj_block[N:2*N,2*N:3*N] = np.identity(N)    # L_23
    
# # use symmetric inter-adjacency matrix
# adj_block += adj_block.T
# mg = mx.MultilayerGraph(list_of_layers=[g1,g2,g3],
#                         inter_adjacency_matrix=adj_block)
mg.set_edges_weights(intra_layer_edges_weight=2,
                     inter_layer_edges_weight=3)

fig = plt.figure(figsize=(15,5))
ax1 = fig.add_subplot(121)
ax1.imshow(mx.adjacency_matrix(mg,weight='weight').todense(),
		  origin='upper',interpolation='nearest',cmap=plt.cm.jet_r)
ax1.set_title('supra adjacency matrix')

ax2 = fig.add_subplot(122)
ax2.axis('off')
ax2.set_title('edge colored network')
pos = mx.get_position(mg,mx.fruchterman_reingold_layout(g1),
					  layer_vertical_shift=1,
					  layer_horizontal_shift=0.0,
					  proj_angle=7)
mx.draw_networkx(mg,pos=pos,ax=ax2,node_size=50,with_labels=False,
				 edge_color=[mg[a][b]['weight'] for a,b in mg.edges()],
				 edge_cmap=plt.cm.jet_r)
plt.show()